#!/usr/bin/env python

# ANTS bot that pursues positions under a learned distribution
#
# Author: Bennett Wineholt

import pickle
import Queue
import sys

from copy import deepcopy
from random import shuffle
from scipy import *
from scipy import ndimage

from ants import *


TRANSITION = 0.125
REMAIN = 0.5

INFLUENCE_SPREAD = 10
FOOD_INFLUENCE_SPREAD = INFLUENCE_SPREAD
# reset in turn 1 # defaults unused
NEED_INFLUENCE_SPREAD = 50
ENEMY_HILL_INFLUENCE_SPREAD = 50

FOOD_INFLUENCE = 20 * sqrt(2) * (1.0/TRANSITION)**FOOD_INFLUENCE_SPREAD
# recalculated in turn 1 # base value expanded at turn 1
NEED_INFLUENCE = 1.01  # unused
ENEMY_HILL_INFLUENCE = 42

class MyBot:
    def __init__(self):
        pass

    def do_setup(self, ants):
        self.log_file = open('feat_log_file', 'wb')
        # pixels never seen
        self.unseen = []
        for row in range(ants.rows):
            for col in range(ants.cols):
                self.unseen.append((row, col))
        self.one_step = array([[0,TRANSITION,0],
                               [TRANSITION,REMAIN,TRANSITION],
                               [0,TRANSITION,0]])
        self.hills = []
        # total influence score of each map pixel
        self.influence_map = zeros((ants.rows, ants.cols))
        # list of locations representing areas needful of more ants
        self.need = []
        # list of locations representative of areas with excess ants
        self.excess = []
        self.WEIGHT =[-0.00010157977670557731, -3.9369582273418096e-05, 0.00010951232303661904, -0.00024465009344827756, 8.109442436456614e-05, -1.5820630758608302e-05] #Replace with regression data
        self.B = 0.0140836675715
        self.hill = []
        self.TURNS_PER = 20 #Hardcoded turns per
        self.DISTANCES = []
        self.CENTERS = []
        self.turn_no = 0


    def do_turn(self, ants):
        # Begin helper method declarations

        def do_move_direction(loc, direction):
            new_loc = ants.destination(loc, direction)
            if (ants.unoccupied(new_loc) and new_loc not in orders):
                ants.issue_order((loc, direction))
                orders[new_loc] = loc
                targets[new_loc] = loc
                return True
            else:
                return False

        def do_move_location(loc, dest):
            directions = ants.direction(loc, dest)
            for direction in directions:
                if do_move_direction(loc, direction):
                    targets[dest] = loc
                    return True
            return False

        # Returns complement of known water on map
        def find_not_water():
            # matrix size of map marking pixels that are not water
            return ((array(ants.map) - WATER) != 0).astype(int)

        def is_pos(x):
            if x > 0:
                return 1.0
            else:
                return 0.0

        # scale need fraction to influence score
        def scale_need(need_fraction):
            return (((10.0 * need_fraction + 1.0) * is_pos(need_fraction)) *
                    (1.0/TRANSITION) ** self.NEED_INFLUENCE_SPREAD)

        # runs feature extractor on presently known state to determine regions
        #   of need and excess for ants
        #
        # sets need to representative list of needful locations
        # sets excess to representative locations of regions with excess ants
        def evaluate_need():
            '''
            Inputs:
                hive location -ants.enemy_hills
                vision information -look in ants
                turn num -looks in self.turn_no
                enemy ants visible ants.enemy_ants
                your total ants ants.my_ants

            '''
            features = getFeatures()
            optimal = []
            for feature_vector in features:
                s = self.B
                for (i,w) in enumerate(self.WEIGHT):
                    s += w*feature_vector[i]
                optimal.append(s)
            s = sum(optimal)
            fractions = [i//len(ants.my_ants()) for i in getNumSq(self.hill,ants.my_ants())]
            self.need_fractions = [(o/s)-fractions[i] for (i,o) in enumerate(optimal)]

            #TODO make need based on need_fractions and CENTER in self

            self.need = [scale_need(need) for need in self.need_fractions]

        '''
            if self.turn_no%5 == 0:
                file = open('influence\data' + str(self.turn_no),'w')
                for i,need in enumerate(self.need):
                    file.write(str(self.CENTERS[i][0]) + ' ' + str(self.CENTERS[i][1]) + ' ' + need)
                file.close()
        '''


        def getFeatures():
            #Feature computations
            f_hills = [0 for i in range(100)]
            for (hill,owner) in ants.enemy_hills():
                f_hills[mapToSquares(self.hill,hill)] += 1.0
            f_vision = getVisionAmount(self.hill,ants.vision)
            f_enemies = getNumSq(self.hill,[tup for (tup,ene) in ants.enemy_ants()])
            f_turns = [self.turn_no//self.TURNS_PER for i in range(100)]
            f_total_ants = [len(ants.my_ants()) for i in range(100)]
            v = [self.DISTANCES, f_turns, f_vision,f_enemies,f_hills,f_total_ants]
            features = [[v[n][m] for n in range(len(v))] for m in range(len(v[0]))]
            return features

        def mapToSquares(hill,location):
            """
            Map the locations by rectangles indicated at the top of
            this function

            """
            nR = 10
            nC = 10
            y0 = hill[0]-ants.rows/2
            x0 = hill[1]-ants.cols/2
            #range generation if decide to not hardcode, to be finished
            # debug
            y_rel = abs(location[0]-y0)%ants.rows
            x_rel = abs(location[1]-x0)%ants.cols
            r = y_rel//(float(ants.rows)/nR)
            c = x_rel//(float(ants.cols)/nC)
            return int(round(r*nC + c))

        def computeCenters(hill,RDIM,CDIM,n):
            (hill_r,hill_c) = hill
            results = []
            r_loc = -RDIM*.5*(1-1.0/n)
            for r in range(n):
                c_loc = -CDIM*.5*(1-1.0/n)
                for c in range(n):
                    results.append((int(r_loc+hill_r)%ants.rows,int(c_loc+hill_c)%ants.cols))
                    c_loc += float(CDIM)/n
            r_loc += float(RDIM)/n
            return results

        def computeDistance(RDIM,CDIM,n):
            results = []
            r_loc = -RDIM*.5*(1-1.0/n)
            for r in range(n):
                c_loc = -CDIM*.5*(1-1.0/n)
                for c in range(n):
                    results.append((r_loc**2 + c_loc**2)**.5)
                    c_loc += float(CDIM)/n
                r_loc += float(RDIM)/n
            return results

        def getNumSq(hill,locs):
            results = [0 for i in range(100)]
            for (r,c) in locs:
                results[mapToSquares(hill,(r,c))] += 1.0

            return results

        def getVisionAmount(hill,vision):
            results = [0 for i in range(100)]
            for r,row_vision in enumerate(vision):
                for c,sq_vision in enumerate(row_vision):
                    if sq_vision:
                        results[mapToSquares(hill,(r,c))] += 1

            return results


        # propagates influence along land by one step
        #
        # inflation corrects for the loss of value given to neighbors since
        # some neighbors are water, 0 sink
        #
        # similarly, to utilize desirable properties of the normalized kernel
        # we re-adjust for the inflated self-influence and maintain balance
        #
        # Requires all scipy array
        #
        # Returns updated map
        def propagate_influence(imap):
            return ((ndimage.correlate(imap, self.one_step)) *
                    self.not_water)
            #return ((ndimage.correlate((imap*self.inflation) +
            #         self.influence_seed, self.one_step) -
            #         REMAIN * (self.inflation-1)) *
            #        self.not_water)

        def process_influence(points, influence, spread):
            imap = zeros((ants.rows, ants.cols))
            for (pt_r, pt_c) in points:
                imap[pt_r][pt_c] = influence
            for i in range(0, spread):
                imap = propagate_influence(imap)
            return imap

        def process_influence_clip(points, influence, spread, minv, maxv):
            imap = zeros((ants.rows, ants.cols))
            for (pt_r, pt_c) in points:
                imap[pt_r][pt_c] = influence
            for i in range(0, spread):
                imap = propagate_influence(imap)
            return clip(imap, minv, maxv)

        def process_influence_list(points, influence_list, spread):
            imap = zeros((ants.rows, ants.cols))
            for i in range(0,len(points)):
                (pt_r, pt_c) = points[i]
                imap[pt_r][pt_c] = influence_list[i]
            for i in range(0, spread):
                imap = propagate_influence(imap)
            return imap

        def process_influence_list_clip(points,
                                        influence_list,
                                        spread,
                                        minv,
                                        maxv):
            imap = zeros((ants.rows, ants.cols))
            for i in range(0,len(points)):
                (pt_r, pt_c) = points[i]
                imap[pt_r][pt_c] = influence_list[i]
            for i in range(0, spread):
                imap = propagate_influence(imap)
            return clip(imap, minv, maxv)

        def process_influence_plus_euclid2(points, influence, spread, origin):
            imap = zeros((ants.rows, ants.cols))
            (o_r, o_c) = origin
            for (pt_r, pt_c) in points:
                imap[pt_r][pt_c] = (influence +
                                    ((o_r-pt_r)**2 + (o_c-pt_c)**2))
            for i in range(0, spread):
                imap = propagate_influence(imap)
            return imap

        def process_influence_plus_euclid2_clip(points,
                                                influence,
                                                spread,
                                                origin,
                                                scale,
                                                minv,
                                                maxv):
            imap = zeros((ants.rows, ants.cols))
            (o_r, o_c) = origin
            for (pt_r, pt_c) in points:
                imap[pt_r][pt_c] = ((influence +
                                     ((o_r-pt_r)**2 + (o_c-pt_c)**2))
                                    / scale)
            for i in range(0, spread):
                imap = propagate_influence(imap)
            return clip(imap, minv, maxv)

        def process_influence_bowl():
            return self.bowl
            #imap = zeros((ants.rows, ants.cols))
            #(o_r, o_c) = self.hill
            #maxv = (FOOD_INFLUENCE * TRANSITION**4)
            #scale = ((((ants.rows/2)**2) + ((ants.cols/2)**2)) / maxv)
            #for pt_r in range(ants.rows):
            #    for pt_c in range(ants.cols):
            #        imap[pt_r][pt_c] = (((o_r-pt_r)**2 + (o_c-pt_c)**2) / scale)
            #return imap

        # update the influence map for a turn with values of each location
        # based on
        # - food
        # - list of locations with need for and excess of ants
        #
        #
        def process_influence_map():
            imap = zeros((ants.rows, ants.cols))
            need_maxv = FOOD_INFLUENCE * (TRANSITION**4)
            imap = (imap +
                    process_influence_list_clip(self.CENTERS,
                                                self.need,
                                                self.NEED_INFLUENCE_SPREAD,
                                                -inf,
                                                need_maxv))
            imap = (imap + process_influence_bowl())
            food_maxv = self.ENEMY_HILL_INFLUENCE * (TRANSITION**10)
            imap = imap + process_influence_clip(ants.food(),
                                                 FOOD_INFLUENCE,
                                                 FOOD_INFLUENCE_SPREAD,
                                                 -inf,
                                                 food_maxv)
            imap = imap + process_influence([hill for (hill, owner) in
                                              ants.enemy_hills()],
                                            self.ENEMY_HILL_INFLUENCE,
                                            self.ENEMY_HILL_INFLUENCE_SPREAD)
            return imap

        # End helper method declarations

        self.turn_no = self.turn_no + 1

        if self.turn_no == 1:
            #Preparation,  to be done once at beginning of game, in appropriate place
            self.hill = ants.my_hills()[0]
            self.DISTANCES = computeDistance(ants.rows,ants.cols,10)
            self.CENTERS = computeCenters(self.hill,ants.rows,ants.cols,10)
            self.NEED_INFLUENCE_SPREAD = max(int(ants.rows/2),
                                        int(ants.cols/2))
            self.ENEMY_HILL_INFLUENCE_SPREAD = max(int(ants.rows/2),
                                              int(ants.cols/2))
            #self.NEED_INFLUENCE = (NEED_INFLUENCE *
            #                       (1.0/TRANSITION)**self.NEED_INFLUENCE_SPREAD)
            self.ENEMY_HILL_INFLUENCE = (ENEMY_HILL_INFLUENCE *
                                         ((1.0/TRANSITION)**
                                          self.ENEMY_HILL_INFLUENCE_SPREAD))
            self.bowl = zeros((ants.rows, ants.cols))
            (o_r, o_c) = self.hill
            bowl_maxv = (FOOD_INFLUENCE * TRANSITION**4)
            bowl_scale = ((((ants.rows/2)**2) + ((ants.cols/2)**2)) / bowl_maxv)
            for pt_r in range(ants.rows):
                for pt_c in range(ants.cols):
                    self.bowl[pt_r][pt_c] = (((o_r-pt_r)**2 + (o_c-pt_c)**2) / bowl_scale)

        # orders to return for each ant
        # dict of ((dest_r, dest_c), (ant_r, ant_c))
        orders = {}
        # destination targets to prevent conflicting pursuit of food
        # dict of ((dest_r, dest_c), (ant_r, ant_c))
        targets = {}

        # update visible map locations
        for loc in self.unseen[:]:
            if ants.visible(loc):
                #self.seen.put(loc)
                self.unseen.remove(loc)

        # update obstacle map, aka what's not water
        self.not_water = find_not_water()

        # run feature extractor to evaulate where ants are needed and not
        # needed
        evaluate_need()

        # influence map accumulates value of map locations
        self.influence_map = process_influence_map()

        # prevent stepping on own hill
        for hill_loc in ants.my_hills():
            orders[hill_loc] = None

        # find food and explore based on collaborative diffusion
        for (ant_r, ant_c) in ants.my_ants():
            pos_dest = []
            for direction in list(AIM.keys()):
                (new_r, new_c) = ants.destination((ant_r, ant_c), direction)
                pos_dest.append((self.influence_map[new_r][new_c],
                                 (new_r, new_c), direction))
            list.sort(pos_dest, reverse=True)
            for (score, dest, direction) in pos_dest:
                if (dest not in targets and
                        (ant_r, ant_c) not in targets.values() and
                        ants.unoccupied(dest)):
                    if do_move_direction((ant_r, ant_c), direction):
                        break
            #orders[(ant_r, ant_c)] = (ant_r, ant_c)

        '''
        Should attack by influence
        # attack hills
        for hill_loc, hill_owner in ants.enemy_hills():
            if hill_loc not in self.hills:
                self.hills.append(hill_loc)
        ant_dist = []
        for hill_loc in self.hills:
            for ant_loc in ants.my_ants():
                if ant_loc not in orders.values():
                    dist = ants.distance(ant_loc, hill_loc)
                    ant_dist.append((dist, ant_loc, hill_loc))
        ant_dist.sort()
        for dist, ant_loc, hill_loc in ant_dist:
            do_move_location(ant_loc, hill_loc)
        '''

        # random move
        # has same effect as exploring unseen without worrying about pathing
        #
        # TODO(bmw): is 2 u-sec enough to terminate?
        for ant_loc in ants.my_ants():
            if ant_loc not in orders.values():
                if (ants.time_remaining() < 2):
                    break
                directions = list(AIM.keys())
                shuffle(directions)
                for direction in directions:
                    dest = ants.destination(ant_loc, direction)
                    if (not dest in targets and
                            ants.unoccupied(dest)):
                        if do_move_location(ant_loc, dest):
                            break

        # unblock own hill
        for hill_loc in ants.my_hills():
            if hill_loc in ants.my_ants() and hill_loc not in orders.values():
                for direction in ('s','e','w','n'):
                    if do_move_direction(hill_loc, direction):
                        break


if __name__ == '__main__':
    # psyco will speed up python a little, but is not needed
    try:
        import psyco
        psyco.full()
    except ImportError:
        pass

    try:
        # if run is passed a class with a do_turn method, it will do the work
        # this is not needed, in which case you will need to write your own
        # parsing function and your own game state class
        Ants.run(MyBot())
    except KeyboardInterrupt:
        print('ctrl-c, leaving ...')
